################################################################################
# AstroImage.py - This module contains a wrapper around the Python Imaging 
#                   Library Image class.
#
"""
TODO: 
"""

__author__ = 'Adrian Price-Whelan <adrn@nyu.edu>'

# Standard library dependencies
import sys, os
import logging

# Third party dependencies
import pyfits as pf
import Image as PIL_Image

# Project Dependencies

def bitpix2PILMode(bitpix):
    if bitpix == 8:
        return "L"
    elif bitpix == 16 or bitpix == 32:
        return "I"
    elif bitpix == -32:
        return "L"
    elif bitpix == -64:
        raise ValueError("Unfortunately 64 Bit images aren't supported by PIL. Complain to them!")
    else:
        raise ValueError("Invalid FITS image.")

def open(fp, hdu=0, mode="r"):
    """ Wraps the PIL Image.open() function to allow opening FITS 
        images.
    """
    if isinstance(fp, pf.core.PrimaryHDU) or isinstance(fp, pf.core.ImageHDU):
        # User supplied an HDU from a FITS file, act accordingly..
        fitsImageData = fp.data
        imMode = bitpix2PILMode(fp.header['BITPIX'])
        return PIL_Image.fromarray(fitsImageData, mode=imMode)
    elif isinstance(fp, "".__class__):
        # User supplied a filename
        filename, ext = os.path.splitext(fp)
        if ext == ".fit" or ext == ".fits":
            hduList = pf.open(fp)
            fitsImageData = hduList[hdu].data
            imMode = bitpix2PILMode(hduList[hdu].header['BITPIX'])
            return PIL_Image.fromarray(fitsImageData, mode=imMode)
        else:
            return PIL_Image.open(fp, mode)

    else:
        return PIL_Image.open(fp, mode)
